from PyQt6.QtWidgets import QInputDialog
import pandas as pd
import os
from datetime import datetime
from openpyxl.styles import PatternFill, Font

class SaveResults:
    def __init__(self, parent):
        self.parent = parent
        self.output_folder = parent.output_folder
        self.excel_file_name = parent.excel_file_name
        self.current_results = parent.current_results

    def save_results(self):
        format, ok = QInputDialog.getItem(
            self.parent, "选择保存格式",
            "请选择保存格式：",
            ["Excel (.xlsx)", "Text (.txt)"],
            0, False
        )
        if not ok:
            return

        if format == "Excel (.xlsx)":
            self.save_to_excel()
        elif format == "Text (.txt)":
            self.save_to_text()

    def save_to_excel(self):
        if not self.current_results:
            self.parent.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            excel_path = os.path.join(self.output_folder, self.excel_file_name)
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            data = []
            separator = {
                "时间戳": "---",
                "图片路径": "---",
                "识别结果": "---"
            }
            for i, (path, text) in enumerate(self.current_results.items()):
                if i > 0:
                    data.append(separator)

                try:
                    rel_path = os.path.relpath(path, self.output_folder)
                except ValueError:
                    rel_path = path

                lines = text.split('\n')
                for line in lines:
                    data.append({
                        "时间戳": timestamp,
                        "图片路径": rel_path,
                        "识别结果": line.strip()
                    })

            new_df = pd.DataFrame(data)

            if os.path.exists(excel_path):
                try:
                    existing_df = pd.read_excel(excel_path, sheet_name='结果')
                    combined_df = pd.concat([existing_df, new_df], ignore_index=True)
                except Exception as e:
                    self.parent.show_warning("读取旧文件失败", f"将创建新文件：{str(e)}")
                    combined_df = new_df
            else:
                combined_df = new_df

            with pd.ExcelWriter(
                    excel_path,
                    engine='openpyxl',
                    mode='w' if not os.path.exists(excel_path) else 'a',
                    if_sheet_exists='replace' if os.path.exists(excel_path) else None
            ) as writer:
                combined_df.to_excel(
                    writer,
                    index=False,
                    sheet_name='结果',
                    header=not os.path.exists(excel_path)
                )

                worksheet = writer.sheets['结果']

                from openpyxl.utils import get_column_letter
                for col in worksheet.columns:
                    max_length = 0
                    column = col[0].column_letter
                    for cell in col:
                        try:
                            cell_value = str(cell.value) if cell.value else ""
                            if len(cell_value) > max_length:
                                max_length = len(cell_value)
                        except:
                            pass
                    adjusted_width = (max_length + 2) * 1.2
                    worksheet.column_dimensions[column].width = adjusted_width

                header_fill = PatternFill(
                    start_color="D3D3D3",
                    end_color="D3D3D3",
                    fill_type="solid"
                )
                header_font = Font(
                    color="000000",
                    bold=True,
                    italic=True
                )

                worksheet.row_dimensions[1].height = 25

            self.parent.show_info("保存成功", f"文件已保存到：\n{excel_path}")

        except PermissionError:
            self.parent.show_error("保存失败", "文件被其他程序占用，请关闭Excel后重试")
        except Exception as e:
            self.parent.show_error("保存失败", f"写入Excel文件时出错：\n{str(e)}")

    def save_to_text(self):
        if not self.current_results:
            self.parent.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            text_path = os.path.join(self.output_folder, self.excel_file_name.replace(".xlsx", ".txt"))
            with open(text_path, "a", encoding="utf-8") as file:
                for path, text in self.current_results.items():
                    file.write(f"=== 图片路径：{path} ===\n")
                    file.write(f"{text}\n\n")

            self.parent.show_info("保存成功", f"文件已保存到：\n{text_path}")

        except Exception as e:
            self.parent.show_error("保存失败", f"写入文本文件时出错：\n{str(e)}")